"""
intensity_dect.py

对音频进行强度计算所用的库.

包含函数:
    normalize():归一化数组, 将给定数组化为每一项与数组最大值的比.
    analyze_wav(): 读取并分析WAV文件.
    do_fft(): 一个用于计算FFT的函数.
    time_domain_intensity():  计算时域上的强度分布，可以选择计算平均功率.
    simple_filter(): 一个简单的滤波器, 用来滤除变化很小的值.
"""



import numpy as np
import time
from scipy.io import wavfile
from scipy.signal import czt
import matplotlib.pyplot as plt


def normalize(y, ctrl_contrast=1.0, ctrl_hide=0.0, range_ctrl = 0, max_val = None):

    """
    归一化数组, 将给定数组化为每一项与数组最大值的比.

    参数:
        y (numpy.ndarray): 要归一化的音频数据.
        ctrl_contrast (float): 对比度, 即每个数据归一化后进行一次幂运算, 该参数作为指数.小
            于1则减小对比度, 大于1则提高对比度.
        ctrl_hide (float): 隐藏值, 小于该值的项不予显示, 或设置为小于, 以提高所有项可见性.
        range_ctrl (bool): 是否启用范围控制, 即将小于0和大于1的项化为0或1.
        max_val (float): 直接给定的最大值, 此种情况下, 忽略y自己的最大值.

    返回:
        y_n (numpy.ndarray): 归一化后的数组y.
    """

    if max_val == None:
        y_max = np.max(np.abs(y))
    else:
        y_max = max_val

    if y_max > 0:
        y_n = y / y_max
    else:
        y_n = y

    y_n = y_n ** ctrl_contrast

    y_n = y_n - ctrl_hide

    if range_ctrl == 1:
        for i in range(len(y_n)):
            if y_n[i] < 0:
                y_n[i] = 0
            if y_n[i] > 1:
                y_n[i] = 1

    return y_n

def analyze_wav(filename, do_fft = 1):

    """
    读取并分析WAV文件.

    参数:
        filename (string): 音频文件地址, 相对路径或绝对路径均可.
        do_fft (bool): 是否进行快速傅里叶变换， 若为1, 则读取完后自动进行一次FFT.

    返回:
        data (numpy.ndarray): 归一化后的数组y.
        sample_rate (int): WAV文件的采样率, 单位Hz.
        freq (numpy.ndarray): 启用FFT后生效, 给出有效的频率值.
        fft_magnitude (numpy.ndarray): 启用FFT后生效, 给出各频率值对应的幅值.
    """

    # 读取WAV文件
    sample_rate, data = wavfile.read(filename)

    # 转换立体声为单声道
    if len(data.shape) > 1:
        data = data.mean(axis=1)

    data = normalize(data).astype(np.float32)

    if do_fft == 1:
        # 计算FFT
        n = len(data)
        fft_result = np.fft.fft(data)
        fft_magnitude = np.abs(fft_result) / n  # 幅度归一化
        fft_magnitude = fft_magnitude[:n//2]    # 取单边频谱

        # 创建频率轴
        freq = np.fft.fftfreq(n, 1/sample_rate)[:n//2]
    else:
        freq = 0
        fft_magnitude = 0

    return data, sample_rate, freq, fft_magnitude

def do_fft(data, sample_rate):

    """
    显而易见, 一个用于计算FFT的函数.

    参数:
        data (numpy.ndarray): 音频数据。
        sample_rate(int): 给出信号的采样率, 单位为Hz, 用于计算时间长度.

    返回:
        data (numpy.ndarray): 归一化后的数组y.
        sample_rate (int): WAV文件的采样率, 单位Hz.
        freq (numpy.ndarray): 有效的频率值.
        fft_magnitude (numpy.ndarray): 各频率值对应的幅值.

    """

    # 计算FFT
    n = len(data)
    fft_result = np.fft.fft(data)
    fft_magnitude = np.abs(fft_result) / n  # 幅度归一化
    fft_magnitude = fft_magnitude[:n//2]    # 取单边频谱

    # 创建频率轴
    freq = np.fft.fftfreq(n, 1/sample_rate)[:n//2]

    return data, sample_rate, freq, fft_magnitude

def time_domain_intensity(data, window_size=None, win_time=None, sample_rate=44100):

    """
    计算时域上的强度分布，可以选择计算平均功率.

    参数:
        data (numpy.ndarray): 音频数据。
        window_size (float, 可选): 用于计算强度的滑动窗口大小(相对于信号长度的比例).
            如果为 None,则计算整个信号的平均功率。
        min_time (int, 可选): 用于直接规定窗口大小, 该项给出时, 忽略window_size, 单位为ms.
            若该项给出, 则需同时给定sample_rate用于计算时间长度.
        sample_rate(int, 可选): 给出信号的采样率, 单位为Hz, 用于计算时间.

    返回:
        average_power (float): 全时间范围内的平均强度.
        intensity(numpy.ndarray): 时域强度分布数据.
        输出为一个二维数组:
        [强度, 时间戳(单位s)]
        为减小开销, 每个窗口内只有一条记录.

    """

    length = len(data)
    if window_size == None and win_time == None:
        # 计算整个信号的平均功率
        total_energy = np.sum(np.abs(data)**2)
        average_power = float(total_energy / length)  # 除以样本数
        return average_power  # 返回平均功率
    else:
        intensity = []

        if win_time != None:
            # 使用滑动窗口(给定时长)计算强度
            window_samples = int(win_time * sample_rate / 2000)  # 窗口大小的一半  \
            window_size = int(win_time * sample_rate / 1000)
            for i in range(0, length, window_size):
                start = max(0, i - window_samples)
                end = min(length, i + window_samples + 1)
                intensity.append([np.sum(np.abs(data[start:end])**2) / (window_size), i/sample_rate])  # 计算窗口内的平均功率
            return intensity
        else:
            # 使用滑动窗口(比例)计算强度
            window_samples = int(window_size * length / 2)  # 窗口大小的一半
            window_size = int(window_size * length)
            for i in range(0, length, window_size):
                start = max(0, i - window_samples)
                end = min(length, i + window_samples + 1)
                intensity.append([np.sum(np.abs(data[start:end])**2) / (window_size), i/sample_rate])  # 计算窗口内的平均功率
            return intensity

def simple_filter(data, normalized=1, delta=0.01):

    """
    一个简单的滤波器, 用来滤除变化很小的值.

    参数:
        data (numpy.ndarray): 带滤波数组.
        normalized (bool): 是否归一化过, 若为0则再归一化一次.
        delta (float): 相对变化率阈值, 前后项之比大于它才会保留, 若不足则设为前一项的值.
    返回:
        res_data (numpy.ndarray): 滤波后数组.
    """

    if normalized == 1:
        res_data = list(data)
        for i in range(1, len(res_data)):
            if res_data[i-1] != 0 and np.abs((res_data[i] - res_data[i-1]) / res_data[i-1]) <= delta:
                res_data[i] = res_data[i-1]
    else:
        res_data = normalize(list(data))
        for i in range(1, len(res_data)):
            if res_data[i-1] != 0 and np.abs((res_data[i] - res_data[i-1]) / res_data[i-1]) <= delta:
                res_data[i] = res_data[i-1]

    return res_data


if __name__ == "__main__":

    # 记录开始时间
    start_time = time.time()

    # 加载音频文件
    file_name = "yinjie_1_1_part.wav"  # 根据需要更新文件路径
    file_data, file_sample_rate, _, _ = analyze_wav(file_name, do_fft=0)
    average_intensity = time_domain_intensity(file_data)
    file_intensity = time_domain_intensity(file_data, win_time=20, sample_rate=file_sample_rate)


    file_intensity_int = list((i[0] for i in file_intensity))
    file_intensity_time = list((i[1] for i in file_intensity))
    file_intensity = simple_filter(file_intensity_int, delta=0.05)


    # 记录结束时间
    end_time = time.time()
    # 计算处理时间
    processing_time = end_time - start_time
    print(f"Processing time: {processing_time:.4f} seconds")

    print(len(file_intensity_int))

    plt.plot(file_intensity_time, file_intensity_int, markersize=8, label='Data Points')

    # 添加标题和标签
    plt.title('Sample Line Plot')
    plt.xlabel('X axis label')
    plt.ylabel('Y axis label')

    # 显示网格
    plt.grid(True)

    # 显示图形
    plt.show()


